"""
LangGraph Agent - 智能课后服务助手
使用状态图管理对话流程和插槽填充
"""
from typing import TypedDict, Annotated, Optional
import operator
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langchain.schema import HumanMessage, SystemMessage, AIMessage

from slots import BookingSlots, IntentType, ServiceType, SlotFiller
from service_manager import ServiceManager


class AgentState(TypedDict):
    """Agent状态定义"""
    messages: Annotated[list, operator.add]  # 对话历史
    slots: BookingSlots  # 插槽信息
    current_slot: Optional[str]  # 当前正在填充的插槽
    user_input: str  # 用户输入
    response: str  # Agent响应
    action_result: Optional[str]  # 操作结果
    finished: bool  # 是否完成


class AfterSchoolServiceAgent:
    """课后服务智能助手"""
    
    def __init__(self, llm: ChatOpenAI):
        self.llm = llm
        self.service_manager = ServiceManager()
        self.slot_filler = SlotFiller()
        self.graph = self._build_graph()
    
    def _build_graph(self) -> StateGraph:
        """构建LangGraph状态图"""
        # 创建状态图
        workflow = StateGraph(AgentState)
        
        # 添加节点
        workflow.add_node("understand_intent", self.understand_intent)
        workflow.add_node("extract_slots", self.extract_slots)
        workflow.add_node("check_slots", self.check_slots)
        workflow.add_node("ask_missing_slot", self.ask_missing_slot)
        workflow.add_node("confirm_action", self.confirm_action)
        workflow.add_node("execute_action", self.execute_action)
        workflow.add_node("generate_response", self.generate_response)
        
        # 设置入口点
        workflow.set_entry_point("understand_intent")
        
        # 添加边
        workflow.add_edge("understand_intent", "extract_slots")
        workflow.add_edge("extract_slots", "check_slots")
        
        # 条件边：检查插槽是否完整
        workflow.add_conditional_edges(
            "check_slots",
            self.should_ask_slot,
            {
                "ask": "ask_missing_slot",
                "confirm": "confirm_action",
                "execute": "execute_action"
            }
        )
        
        workflow.add_edge("ask_missing_slot", END)
        workflow.add_edge("confirm_action", END)
        workflow.add_edge("execute_action", "generate_response")
        workflow.add_edge("generate_response", END)
        
        return workflow.compile()
    
    def understand_intent(self, state: AgentState) -> AgentState:
        """理解用户意图"""
        user_input = state["user_input"]
        
        # 如果已有意图且正在填充插槽，保持当前意图
        if state["slots"].intent and state["current_slot"]:
            return state
        
        # 提取意图
        intent = self.slot_filler.extract_intent(user_input)
        state["slots"].intent = intent
        
        # 使用LLM增强意图理解
        if intent == IntentType.UNKNOWN:
            system_prompt = """你是一个课后服务助手。分析用户意图并分类为以下之一：
- book: 用户想要预定课后服务
- cancel: 用户想要取消预定
- query: 用户想要查询自己的预定
- list_all: 用户想要查看所有预定
- unknown: 无法确定意图

只返回意图类型，不要其他内容。"""
            
            messages = [
                SystemMessage(content=system_prompt),
                HumanMessage(content=user_input)
            ]
            
            response = self.llm.invoke(messages)
            intent_str = response.content.strip().lower()
            
            try:
                state["slots"].intent = IntentType(intent_str)
            except ValueError:
                state["slots"].intent = IntentType.UNKNOWN
        
        return state
    
    def extract_slots(self, state: AgentState) -> AgentState:
        """从用户输入中提取插槽信息"""
        user_input = state["user_input"]
        slots = state["slots"]
        
        # 提取服务类型
        if not slots.service_type:
            service_type = self.slot_filler.extract_service_type(user_input)
            if service_type:
                slots.service_type = service_type
        
        # 提取学生姓名
        if not slots.student_name:
            student_name = self.slot_filler.extract_student_name(user_input, slots)
            if student_name:
                slots.student_name = student_name
        
        # 使用LLM提取其他信息
        if slots.intent in [IntentType.BOOK, IntentType.QUERY]:
            extracted = self._llm_extract_slots(user_input, slots)
            
            # 更新插槽
            if "date" in extracted and not slots.date:
                slots.date = extracted["date"]
            if "time_slot" in extracted and not slots.time_slot:
                slots.time_slot = extracted["time_slot"]
            if "teacher_preference" in extracted and not slots.teacher_preference:
                slots.teacher_preference = extracted["teacher_preference"]
            if "booking_id" in extracted and not slots.booking_id:
                slots.booking_id = extracted["booking_id"]
        
        state["slots"] = slots
        return state
    
    def _llm_extract_slots(self, user_input: str, slots: BookingSlots) -> dict:
        """使用LLM提取插槽信息"""
        system_prompt = f"""从用户输入中提取以下信息（JSON格式）：
- date: 日期（YYYY-MM-DD格式）
- time_slot: 时间段（HH:MM-HH:MM格式）
- teacher_preference: 教师偏好
- booking_id: 预定编号
- student_name: 学生姓名

当前意图: {slots.intent.value if slots.intent else 'unknown'}

只返回JSON对象，包含提取到的信息。如果某个信息未提及，不要包含该字段。
示例: {{"date": "2024-01-15", "time_slot": "14:00-16:00"}}"""
        
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_input)
        ]
        
        try:
            response = self.llm.invoke(messages)
            content = response.content.strip()
            
            # 提取JSON
            import json
            if "```json" in content:
                content = content.split("```json")[1].split("```")[0].strip()
            elif "```" in content:
                content = content.split("```")[1].split("```")[0].strip()
            
            return json.loads(content)
        except Exception as e:
            print(f"LLM提取失败: {e}")
            return {}
    
    def check_slots(self, state: AgentState) -> AgentState:
        """检查插槽完整性"""
        slots = state["slots"]
        missing = slots.get_missing_slots()
        state["slots"].missing_slots = missing
        return state
    
    def should_ask_slot(self, state: AgentState) -> str:
        """决定下一步动作"""
        slots = state["slots"]
        
        # 意图未知
        if slots.intent == IntentType.UNKNOWN:
            return "ask"
        
        # 不需要插槽的操作
        if slots.intent == IntentType.LIST_ALL:
            return "execute"
        
        # 检查是否有缺失的插槽
        if not slots.is_complete():
            return "ask"
        
        # 如果是预定且未确认
        if slots.intent == IntentType.BOOK and not slots.confirmed:
            return "confirm"
        
        return "execute"
    
    def ask_missing_slot(self, state: AgentState) -> AgentState:
        """询问缺失的插槽"""
        slots = state["slots"]
        
        if slots.intent == IntentType.UNKNOWN:
            state["response"] = """您好！我是课后服务助手。我可以帮您：

1. 📝 预定课后服务
2. ❌ 取消预定
3. 🔍 查询您的预定
4. 📋 查看所有预定

请告诉我您需要什么帮助？"""
            state["finished"] = False
            return state
        
        missing_slots = slots.get_missing_slots()
        
        if missing_slots:
            # 询问第一个缺失的插槽
            next_slot = missing_slots[0]
            state["current_slot"] = next_slot
            state["response"] = slots.get_slot_prompt_message(next_slot)
            state["finished"] = False
        
        return state
    
    def confirm_action(self, state: AgentState) -> AgentState:
        """确认操作"""
        slots = state["slots"]
        
        confirm_msg = f"""
请确认以下预定信息：

{slots}

确认无误请回复"确认"或"是"，需要修改请告诉我。
        """.strip()
        
        # 检查用户是否确认
        user_input_lower = state["user_input"].lower()
        if any(word in user_input_lower for word in ["确认", "是", "对", "没错", "可以", "ok", "yes"]):
            slots.confirmed = True
            state["slots"] = slots
            return self.execute_action(state)
        
        state["response"] = confirm_msg
        state["finished"] = False
        return state
    
    def execute_action(self, state: AgentState) -> AgentState:
        """执行操作"""
        slots = state["slots"]
        intent = slots.intent
        
        if intent == IntentType.BOOK:
            success, message, booking = self.service_manager.create_booking(slots)
            state["action_result"] = message
            state["finished"] = True
            
        elif intent == IntentType.CANCEL:
            if slots.booking_id:
                success, message = self.service_manager.cancel_booking(slots.booking_id)
            else:
                success, message = False, "缺少预定编号"
            state["action_result"] = message
            state["finished"] = True
            
        elif intent == IntentType.QUERY:
            if slots.student_name:
                success, message, bookings = self.service_manager.query_bookings(slots.student_name)
            else:
                success, message = False, "缺少学生姓名"
            state["action_result"] = message
            state["finished"] = True
            
        elif intent == IntentType.LIST_ALL:
            success, message, bookings = self.service_manager.list_all_bookings()
            state["action_result"] = message
            state["finished"] = True
        
        return state
    
    def generate_response(self, state: AgentState) -> AgentState:
        """生成最终响应"""
        if state.get("action_result"):
            state["response"] = state["action_result"]
        
        return state
    
    def chat(self, user_input: str, state: Optional[AgentState] = None) -> tuple[str, AgentState]:
        """
        处理用户输入
        
        Args:
            user_input: 用户输入
            state: 当前状态（用于多轮对话）
        
        Returns:
            (响应, 新状态)
        """
        if state is None:
            state = {
                "messages": [],
                "slots": BookingSlots(),
                "current_slot": None,
                "user_input": user_input,
                "response": "",
                "action_result": None,
                "finished": False
            }
        else:
            state["user_input"] = user_input
        
        # 添加用户消息
        state["messages"].append(HumanMessage(content=user_input))
        
        # 运行图
        result = self.graph.invoke(state)
        
        # 添加助手响应
        response = result.get("response", "抱歉，我没有理解您的意思。")
        result["messages"].append(AIMessage(content=response))
        
        return response, result
